# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import collections
import logging
import json
from typing import List, Union

from tornado import gen
import zlib

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


def parse_channel(channel):
  market, symbol, event_type, arg = channel.split('.')
  return symbol, event_type


TopicPatterns = collections.namedtuple('TopicPatterns',
                                       ['trade', 'depth', 'delta_depth', 'topic_event'])


class HuobiSubscriberBase(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _decode(self, ws_msg):
    return json.loads(zlib.decompress(bytes(ws_msg), 15 + 32))

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('%s_%s%s' %
                                            (self._topics.depth, group_id, worker_id))[0]
    topic_delta_depth = topic_map.search_topic_id(
        '%s_%s%s' % (self._topics.delta_depth, group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('%s_%s%s' %
                                            (self._topics.trade, group_id, worker_id))[0]

    self._publisher = {
        'depth': queue.get_writer(topic_id=topic_depth[1]),
        'mbp': queue.get_writer(topic_id=topic_delta_depth[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1])
    }

    self._enable_topic_event_publisher(queue,
                                       self._topics.topic_event, [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    yield self._send_subscribe_request()

  @gen.coroutine
  def _send_subscribe_request(self):
    for symbol in self._symbol:
      self._logger.debug('Subscribing %s' % symbol)
      yield self._ws.write_message(
          json.dumps({
              'sub': 'market.%s.depth.step0' % symbol, 'id': '%s.depth' % symbol
          }))
      yield self._ws.write_message(
          json.dumps({
              'sub': 'market.%s.mbp.150' % symbol, 'id': '%s.depth' % symbol
          }))
      yield self._ws.write_message(
          json.dumps({
              'sub': 'market.%s.trade.detail' % symbol, 'id': '%s.trade' % symbol
          }))
      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  @gen.coroutine
  def pong(self, timestamp):
    payload = json.dumps({'pong': timestamp})
    yield self._ws.write_message(payload)

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return
    if 'ping' in message:
      ping_timestamp = message['ping']
      self._ioloop.add_callback(lambda: self.pong(ping_timestamp))
      return
    if 'subbed' in message:
      # TODO(inkyu): Archive me
      return

    channel = message.get('ch', '')
    try:
      symbol, event_type = parse_channel(channel)
      self._publisher[event_type].write(timestamp, message)
    except (KeyError, ValueError):
      self._logger.error('Unknown channel: %s' % message)


class HuobiSubscriber(HuobiSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    self._topics = TopicPatterns(
        trade='huobi_ws_trade',
        depth='huobi_ws_depth',
        delta_depth='huobi_ws_delta_depth',
        topic_event='huobi_ws_topic_event',
    )
    super().__init__(queue, group_id, worker_id, symbol, close_callback, logger)
    self._register_url('wss://api.huobi.pro/ws')
